"""
ACME商店语音助手系统 - 完全独立版本
不使用agents模块，完全本地实现所有功能
"""

import os
import PyPDF2
import faiss
from sentence_transformers import SentenceTransformer
import asyncio
from typing import List
from datetime import datetime

print("正在初始化本地AI助手系统...")

# 初始化本地嵌入模型
print("正在加载嵌入模型...")
embedding_model = SentenceTransformer(r'D:\ideaSpace\MyPython\models\all-MiniLM-L6-v2')

class LocalVectorStore:
    """本地向量存储实现"""

    def __init__(self, store_name: str):
        self.store_name = store_name
        self.index = None
        self.documents = []
        self.metadata = []
        self.dimension = 384

    def create(self):
        """创建向量索引"""
        self.index = faiss.IndexFlatL2(self.dimension)
        return {
            "id": f"local_{self.store_name}",
            "name": self.store_name,
            "created_at": datetime.now().isoformat(),
            "file_count": len(self.documents)
        }

    def add_document(self, text: str, metadata: dict = None):
        """添加文档到向量存储"""
        if metadata is None:
            metadata = {}

        embedding = embedding_model.encode(text)
        embedding = embedding.reshape(1, -1)

        if self.index is None:
            self.index = faiss.IndexFlatL2(self.dimension)

        self.index.add(embedding.astype('float32'))
        self.documents.append(text)
        self.metadata.append(metadata)

    def search(self, query: str, max_results: int = 3):
        """搜索相似文档"""
        if self.index is None or len(self.documents) == 0:
            return self._get_fallback_results(query)

        try:
            query_embedding = embedding_model.encode(query)
            query_embedding = query_embedding.reshape(1, -1).astype('float32')

            distances, indices = self.index.search(query_embedding, max_results)

            results = []
            for i, idx in enumerate(indices[0]):
                if idx < len(self.documents):
                    results.append({
                        "content": self.documents[idx],
                        "metadata": self.metadata[idx],
                        "score": float(distances[0][i])
                    })

            return results if results else self._get_fallback_results(query)
        except Exception as e:
            print(f"搜索时出错: {e}")
            return self._get_fallback_results(query)

    def _get_fallback_results(self, query: str):
        """获取备用搜索结果"""
        fallback_data = {
            "account": [{
                "content": "用户账户信息：用户ID 1234567890 的余额为£72.50，会员状态是Gold Executive",
                "metadata": {"source": "fallback", "type": "account"},
                "score": 0.9
            }],
            "dynamite": [{
                "content": "ACME炸药分配器：输入尺寸10cm，输出速度每秒5个，安全认证级别A+",
                "metadata": {"source": "fallback", "type": "product"},
                "score": 0.9
            }],
            "duck": [{
                "content": "鸭子狩猎装备：包括专业伪装服、电子诱饵系统、高精度猎枪和配件",
                "metadata": {"source": "fallback", "type": "product"},
                "score": 0.9
            }],
            "default": [{
                "content": f"关于'{query}'的信息：ACME公司提供各种高质量产品和服务",
                "metadata": {"source": "fallback", "type": "general"},
                "score": 0.8
            }]
        }

        query_lower = query.lower()
        if "account" in query_lower or "balance" in query_lower:
            return fallback_data["account"]
        elif "dynamite" in query_lower or "dispenser" in query_lower:
            return fallback_data["dynamite"]
        elif "duck" in query_lower or "hunt" in query_lower:
            return fallback_data["duck"]
        else:
            return fallback_data["default"]

# 模拟函数工具装饰器
def function_tool(func):
    """模拟function_tool装饰器"""
    func.is_tool = True
    return func

class Tool:
    """工具基类"""
    def __init__(self, name, description):
        self.name = name
        self.description = description

    async def execute(self, *args, **kwargs):
        raise NotImplementedError

class WebSearchTool(Tool):
    """模拟网络搜索工具"""
    def __init__(self):
        super().__init__("web_search", "搜索网络信息")

    async def execute(self, query: str):
        """执行网络搜索"""
        search_results = {
            "duck hunting gear": [
                {
                    "title": "2024年鸭子狩猎装备趋势",
                    "snippet": "智能诱饵和热成像技术成为新趋势",
                    "url": "https://example.com/trends"
                }
            ],
            "default": [
                {
                    "title": f"{query} 的搜索结果",
                    "snippet": f"找到关于'{query}'的相关信息",
                    "url": "https://example.com/search"
                }
            ]
        }

        return search_results.get(query.lower(), search_results["default"])

class FileSearchTool(Tool):
    """文件搜索工具"""
    def __init__(self, vector_store, max_num_results=3):
        super().__init__("file_search", "搜索文件内容")
        self.vector_store = vector_store
        self.max_num_results = max_num_results

    async def execute(self, query: str):
        """执行文件搜索"""
        results = self.vector_store.search(query, self.max_num_results)
        return [{"content": r["content"], "metadata": r["metadata"]} for r in results]

class Agent:
    """代理类"""
    def __init__(self, name: str, instructions: str, tools: List[Tool] = None):
        self.name = name
        self.instructions = instructions
        self.tools = tools or []

    async def process_query(self, query: str):
        """处理查询"""
        print(f"{self.name} 正在处理查询: {query}")

        # 简单的意图识别
        query_lower = query.lower()

        # 检查是否需要使用工具
        for tool in self.tools:
            if hasattr(tool, 'name'):
                if tool.name == "get_account_info" and ("account" in query_lower or "balance" in query_lower):
                    result = await tool.execute("1234567890")
                    return self._format_account_response(result)
                elif tool.name == "file_search" and ("product" in query_lower or "dynamite" in query_lower or "duck" in query_lower):
                    result = await tool.execute(query)
                    return self._format_product_response(result, query)
                elif tool.name == "web_search":
                    result = await tool.execute(query)
                    return self._format_web_response(result)

        # 默认响应
        return self._format_default_response(query)

    def _format_account_response(self, account_info):
        """格式化账户响应"""
        return f"账户信息：{account_info['name']}，余额：{account_info['account_balance']}，会员状态：{account_info['membership_status']}"

    def _format_product_response(self, results, query):
        """格式化产品响应"""
        if not results:
            return f"抱歉，没有找到关于'{query}'的产品信息。"

        response = f"关于'{query}'的产品信息：\n"
        for i, result in enumerate(results[:3], 1):
            response += f"{i}. {result['content']}\n"
        return response

    def _format_web_response(self, results):
        """格式化网络响应"""
        if not results:
            return "抱歉，没有找到相关的网络信息。"

        response = "网络搜索结果：\n"
        for i, result in enumerate(results[:3], 1):
            response += f"{i}. {result['title']}: {result['snippet']}\n"
        return response

    def _format_default_response(self, query):
        """格式化默认响应"""
        responses = {
            "greeting": "您好！我是ACME商店助手，请问有什么可以帮您？",
            "thanks": "不客气！很高兴为您服务。",
            "default": f"您好！关于'{query}'，我可以帮您查询账户信息、产品详情或进行网络搜索。"
        }

        query_lower = query.lower()
        if any(word in query_lower for word in ["hello", "hi", "hey", "你好"]):
            return responses["greeting"]
        elif any(word in query_lower for word in ["thank", "thanks", "谢谢"]):
            return responses["thanks"]
        else:
            return responses["default"]

class Runner:
    """运行器类"""
    @staticmethod
    async def run(agent, query):
        """运行代理"""
        class Result:
            def __init__(self, output):
                self.final_output = output

        try:
            response = await agent.process_query(query)
            return Result(response)
        except Exception as e:
            print(f"代理执行错误: {e}")
            return Result(f"抱歉，处理您的查询时出现错误: {e}")

# --- 工具函数 ---

@function_tool
async def get_account_info(user_id: str) -> dict:
    """
    获取用户账户信息
    """
    return {
        "user_id": user_id,
        "name": "Bugs Bunny",
        "account_balance": "£72.50",
        "membership_status": "Gold Executive"
    }

# --- 辅助函数 ---

def extract_text_from_pdf(pdf_path: str):
    """从PDF文件中提取文本"""
    text = ""
    try:
        if os.path.exists(pdf_path):
            with open(pdf_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    text += page.extract_text() + "\n"
            print(f"成功从PDF提取文本，长度: {len(text)} 字符")
        else:
            print(f"PDF文件不存在: {pdf_path}")
            text = """
            ACME产品目录
            炸药分配器: 输入尺寸10cm, 输出速度每秒5个
            鸭子狩猎装备: 专业级装备套装
            火箭滑板车: 最高速度100km/h
            """
    except Exception as e:
        print(f"提取PDF文本时出错: {e}")
        text = "ACME公司提供各种高质量产品和服务"

    return text

def create_vector_store(store_name: str, pdf_path: str = None) -> LocalVectorStore:
    """
    创建本地向量存储库
    """
    vector_store = LocalVectorStore(store_name)
    vector_store.create()

    if pdf_path:
        print(f"正在处理PDF文件: {pdf_path}")
        text = extract_text_from_pdf(pdf_path)

        # 将文本分割成chunks
        chunks = [chunk for chunk in text.split('\n\n') if chunk.strip()]
        for i, chunk in enumerate(chunks):
            vector_store.add_document(chunk, {"source": pdf_path, "chunk_id": i})

        print(f"成功添加 {len(chunks)} 个文本块到向量存储")
    else:
        # 添加默认数据
        default_data = [
            "ACME炸药分配器具有10cm输入尺寸和每秒5个输出速度",
            "鸭子狩猎装备包括专业伪装服和电子诱饵系统",
            "用户账户查询请提供用户ID",
            "ACME产品质量保证，所有产品都有1年保修"
        ]
        for i, data in enumerate(default_data):
            vector_store.add_document(data, {"source": "default", "chunk_id": i})
        print("已添加默认数据到向量存储")

    return vector_store

# --- 初始化系统 ---

print("创建向量存储...")
vector_store = create_vector_store(
    "ACME知识库",
    "D:/tmp/acme_product_catalogue.pdf"
)

print("初始化工具...")
web_search_tool = WebSearchTool()
file_search_tool = FileSearchTool(vector_store)
account_tool = type('AccountTool', (), {
    'name': 'get_account_info',
    'execute': get_account_info
})()

print("初始化代理...")
# 搜索代理
search_agent = Agent(
    name="SearchAgent",
    instructions="处理需要网络搜索的查询",
    tools=[web_search_tool]
)

# 知识代理
knowledge_agent = Agent(
    name="KnowledgeAgent",
    instructions="处理产品相关查询",
    tools=[file_search_tool]
)

# 账户代理
account_agent = Agent(
    name="AccountAgent",
    instructions="处理账户相关查询",
    tools=[account_tool]
)

# 分流代理
triage_agent = Agent(
    name="Assistant",
    instructions="欢迎用户并路由到合适的专业代理",
    tools=[web_search_tool, file_search_tool, account_tool]
)

# --- 测试函数 ---

async def test_queries():
    """
    测试不同查询的响应
    """
    examples = [
        "What's my ACME account balance? My user ID is 1234567890",
        "How big is the input and how fast is the output of the dynamite dispenser?",
        "What duck hunting gear products do you have?",
        "Hello, how are you?",
        "Thank you for your help!"
    ]

    print("\n开始测试查询...")
    print("=" * 60)

    for query in examples:
        print(f"\n测试查询: {query}")
        print("-" * 50)

        # 根据查询内容选择代理
        query_lower = query.lower()
        if "account" in query_lower or "balance" in query_lower:
            agent = account_agent
        elif "dynamite" in query_lower or "product" in query_lower:
            agent = knowledge_agent
        elif "duck" in query_lower or "hunt" in query_lower:
            agent = knowledge_agent
        else:
            agent = triage_agent

        result = await Runner.run(agent, query)
        print(f"用户: {query}")
        print(f"助手: {result.final_output}")
        print("---")

# --- 交互式聊天 ---

async def interactive_chat():
    """
    交互式聊天模式
    """
    print("\n交互式聊天模式已启动")
    print("输入 'exit' 或 'quit' 退出程序")
    print("=" * 50)

    while True:
        try:
            user_input = input("\n您: ").strip()

            if user_input.lower() in ['exit', 'quit', 'bye']:
                print("助手: 再见！感谢使用ACME商店助手。")
                break

            if not user_input:
                continue

            # 选择代理
            query_lower = user_input.lower()
            if "account" in query_lower or "balance" in query_lower:
                agent = account_agent
            elif "dynamite" in query_lower or "duck" in query_lower or "product" in query_lower:
                agent = knowledge_agent
            elif "search" in query_lower or "web" in query_lower:
                agent = search_agent
            else:
                agent = triage_agent

            print("助手: 思考中...", end="\r")
            result = await Runner.run(agent, user_input)
            print(f"助手: {result.final_output}")

        except KeyboardInterrupt:
            print("\n\n助手: 再见！")
            break
        except Exception as e:
            print(f"\n助手: 抱歉，处理时出现错误: {e}")

# --- 主执行块 ---

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='ACME本地AI助手')
    parser.add_argument('--test', action='store_true', help='运行测试查询')
    parser.add_argument('--chat', action='store_true', help='启动交互式聊天')
    args = parser.parse_args()

    print("ACME本地AI助手系统初始化完成!")
    print("=" * 50)

    if args.test:
        asyncio.run(test_queries())
    elif args.chat:
        asyncio.run(interactive_chat())
    else:
        # 默认运行测试然后进入聊天模式
        asyncio.run(test_queries())
        print("\n" + "=" * 50)
        asyncio.run(interactive_chat())